"""
文件名: rate_limiter.py
描述: API速率限制器

本模块包含以下主要功能：
1. 基于Redis的分布式速率限制
2. 多种限流算法（固定窗口、滑动窗口、令牌桶）
3. 用户级别和IP级别限流
4. YouTube API配额管理

依赖模块:
   - redis: Redis客户端
   - time: 时间处理
   - core: 配置和日志模块

使用示例:
   >>> @rate_limit(max_requests=100, window=60)
   >>> async def api_function():
   >>>     return "result"

注意事项:
   - 分布式环境下要使用Redis存储
   - 限流策略要根据业务需求调整
   - YouTube API配额要特别管理
   - 要考虑突发流量的处理

作者: AI助手
创建日期: 2024-06-27
最后修改: 2024-06-27
版本: 1.0.0
"""

import time
import json
import asyncio
from datetime import datetime, timedelta
from typing import Dict, Any, Optional, List, Callable
from functools import wraps
from enum import Enum

from fastapi import HTTPException, Request, status
from src.core.redis import get_redis_manager
from src.core.logger import get_logger
from src.core.config import get_settings

settings = get_settings()
logger = get_logger(__name__)


class RateLimitAlgorithm(Enum):
    """限流算法类型"""
    FIXED_WINDOW = "fixed_window"
    SLIDING_WINDOW = "sliding_window"
    TOKEN_BUCKET = "token_bucket"
    LEAKY_BUCKET = "leaky_bucket"


class RateLimitScope(Enum):
    """限流范围类型"""
    GLOBAL = "global"
    USER = "user"
    IP = "ip"
    ENDPOINT = "endpoint"


class RateLimiter:
    """通用速率限制器"""
    
    def __init__(self):
        self.redis = None
        self.prefix = f"{settings.cache_prefix}:rate_limit"
    
    async def initialize(self):
        """初始化限流器"""
        try:
            redis_manager = await get_redis_manager()
            self.redis = redis_manager.client
            logger.info("速率限制器初始化成功")
        except Exception as e:
            logger.error(f"速率限制器初始化失败: {e}")
            raise
    
    def _get_key(self, scope: RateLimitScope, identifier: str, window: int = None) -> str:
        """生成限流键"""
        timestamp = int(time.time())
        
        if window:
            # 固定窗口：按窗口时间分组
            window_start = timestamp // window * window
            return f"{self.prefix}:{scope.value}:{identifier}:{window_start}"
        else:
            # 其他算法使用基础键
            return f"{self.prefix}:{scope.value}:{identifier}"
    
    async def check_rate_limit(
        self,
        scope: RateLimitScope,
        identifier: str,
        max_requests: int,
        window_seconds: int,
        algorithm: RateLimitAlgorithm = RateLimitAlgorithm.FIXED_WINDOW
    ) -> Dict[str, Any]:
        """
        检查速率限制
        
        Args:
            scope: 限流范围
            identifier: 标识符（用户ID、IP等）
            max_requests: 最大请求数
            window_seconds: 时间窗口（秒）
            algorithm: 限流算法
            
        Returns:
            Dict[str, Any]: 限流检查结果
        """
        if not self.redis:
            # Redis不可用时，允许请求通过
            return {
                "allowed": True,
                "remaining": max_requests,
                "reset_at": time.time() + window_seconds,
                "retry_after": 0
            }
        
        try:
            if algorithm == RateLimitAlgorithm.FIXED_WINDOW:
                return await self._fixed_window_check(
                    scope, identifier, max_requests, window_seconds
                )
            elif algorithm == RateLimitAlgorithm.SLIDING_WINDOW:
                return await self._sliding_window_check(
                    scope, identifier, max_requests, window_seconds
                )
            elif algorithm == RateLimitAlgorithm.TOKEN_BUCKET:
                return await self._token_bucket_check(
                    scope, identifier, max_requests, window_seconds
                )
            else:
                # 默认使用固定窗口
                return await self._fixed_window_check(
                    scope, identifier, max_requests, window_seconds
                )
                
        except Exception as e:
            logger.error(f"速率限制检查失败: {e}")
            # 出错时允许请求通过
            return {
                "allowed": True,
                "remaining": max_requests,
                "reset_at": time.time() + window_seconds,
                "retry_after": 0
            }
    
    async def _fixed_window_check(
        self, scope: RateLimitScope, identifier: str, max_requests: int, window_seconds: int
    ) -> Dict[str, Any]:
        """固定窗口限流检查"""
        key = self._get_key(scope, identifier, window_seconds)
        
        pipe = self.redis.pipeline()
        pipe.incr(key)
        pipe.expire(key, window_seconds)
        results = await pipe.execute()
        
        current_requests = results[0]
        
        remaining = max(0, max_requests - current_requests)
        allowed = current_requests <= max_requests
        
        # 计算重置时间
        current_window_start = int(time.time()) // window_seconds * window_seconds
        reset_at = current_window_start + window_seconds
        
        return {
            "allowed": allowed,
            "remaining": remaining,
            "reset_at": reset_at,
            "retry_after": reset_at - time.time() if not allowed else 0,
            "current_requests": current_requests
        }
    
    async def _sliding_window_check(
        self, scope: RateLimitScope, identifier: str, max_requests: int, window_seconds: int
    ) -> Dict[str, Any]:
        """滑动窗口限流检查"""
        key = self._get_key(scope, identifier)
        now = time.time()
        cutoff = now - window_seconds
        
        pipe = self.redis.pipeline()
        # 移除过期的记录
        pipe.zremrangebyscore(key, 0, cutoff)
        # 添加当前请求
        pipe.zadd(key, {str(now): now})
        # 获取窗口内的请求数
        pipe.zcard(key)
        # 设置过期时间
        pipe.expire(key, window_seconds)
        
        results = await pipe.execute()
        current_requests = results[2]
        
        remaining = max(0, max_requests - current_requests)
        allowed = current_requests <= max_requests
        
        return {
            "allowed": allowed,
            "remaining": remaining,
            "reset_at": now + window_seconds,
            "retry_after": 1 if not allowed else 0,  # 滑动窗口建议1秒后重试
            "current_requests": current_requests
        }
    
    async def _token_bucket_check(
        self, scope: RateLimitScope, identifier: str, capacity: int, refill_rate: int
    ) -> Dict[str, Any]:
        """令牌桶限流检查"""
        key = self._get_key(scope, identifier)
        now = time.time()
        
        # 获取当前桶状态
        bucket_data = await self.redis.get(key)
        
        if bucket_data:
            bucket = json.loads(bucket_data)
            last_refill = bucket['last_refill']
            tokens = bucket['tokens']
        else:
            # 初始化桶
            last_refill = now
            tokens = capacity
        
        # 计算需要添加的令牌数
        time_passed = now - last_refill
        tokens_to_add = int(time_passed * refill_rate)
        tokens = min(capacity, tokens + tokens_to_add)
        
        # 检查是否有足够令牌
        if tokens >= 1:
            tokens -= 1
            allowed = True
        else:
            allowed = False
        
        # 更新桶状态
        new_bucket = {
            'tokens': tokens,
            'last_refill': now,
            'capacity': capacity,
            'refill_rate': refill_rate
        }
        
        await self.redis.setex(key, 3600, json.dumps(new_bucket))
        
        return {
            "allowed": allowed,
            "remaining": tokens,
            "reset_at": now + (capacity - tokens) / refill_rate if refill_rate > 0 else now + 60,
            "retry_after": (1 - tokens) / refill_rate if not allowed and refill_rate > 0 else 1,
            "current_tokens": tokens
        }


class YouTubeAPIQuotaManager:
    """YouTube API配额管理器"""
    
    def __init__(self):
        self.redis = None
        self.prefix = f"{settings.cache_prefix}:youtube_quota"
        self.daily_quota = 10000  # YouTube API默认每日配额
        self.cost_map = {
            'search': 100,
            'videos': 1,
            'channels': 1,
            'subscriptions': 1,
            'playlistItems': 1,
        }
    
    async def initialize(self):
        """初始化配额管理器"""
        try:
            redis_manager = await get_redis_manager()
            self.redis = redis_manager.client
            logger.info("YouTube API配额管理器初始化成功")
        except Exception as e:
            logger.error(f"YouTube API配额管理器初始化失败: {e}")
            raise
    
    async def check_quota(self, operation: str, user_id: str = None) -> Dict[str, Any]:
        """
        检查API配额
        
        Args:
            operation: API操作类型
            user_id: 用户ID（可选）
            
        Returns:
            Dict[str, Any]: 配额检查结果
        """
        if not self.redis:
            return {"allowed": True, "remaining": self.daily_quota}
        
        cost = self.cost_map.get(operation, 1)
        
        # 全局配额检查
        global_key = f"{self.prefix}:global:{self._get_date_key()}"
        global_used = await self.redis.get(global_key) or 0
        global_used = int(global_used)
        
        if global_used + cost > self.daily_quota:
            return {
                "allowed": False,
                "remaining": max(0, self.daily_quota - global_used),
                "reason": "全局配额不足",
                "reset_at": self._get_next_reset_time()
            }
        
        # 用户配额检查（如果提供用户ID）
        if user_id:
            user_quota = 1000  # 每用户每日配额
            user_key = f"{self.prefix}:user:{user_id}:{self._get_date_key()}"
            user_used = await self.redis.get(user_key) or 0
            user_used = int(user_used)
            
            if user_used + cost > user_quota:
                return {
                    "allowed": False,
                    "remaining": max(0, user_quota - user_used),
                    "reason": "用户配额不足",
                    "reset_at": self._get_next_reset_time()
                }
        
        return {
            "allowed": True,
            "remaining": self.daily_quota - global_used,
            "cost": cost
        }
    
    async def consume_quota(self, operation: str, user_id: str = None) -> bool:
        """
        消费API配额
        
        Args:
            operation: API操作类型
            user_id: 用户ID（可选）
            
        Returns:
            bool: 是否成功消费
        """
        if not self.redis:
            return True
        
        cost = self.cost_map.get(operation, 1)
        
        # 检查配额
        quota_check = await self.check_quota(operation, user_id)
        if not quota_check["allowed"]:
            return False
        
        try:
            # 消费全局配额
            global_key = f"{self.prefix}:global:{self._get_date_key()}"
            pipe = self.redis.pipeline()
            pipe.incrby(global_key, cost)
            pipe.expire(global_key, 86400)  # 24小时过期
            
            # 消费用户配额
            if user_id:
                user_key = f"{self.prefix}:user:{user_id}:{self._get_date_key()}"
                pipe.incrby(user_key, cost)
                pipe.expire(user_key, 86400)
            
            await pipe.execute()
            
            logger.info(f"YouTube API配额消费", operation=operation, cost=cost, user_id=user_id)
            return True
            
        except Exception as e:
            logger.error(f"配额消费失败: {e}")
            return False
    
    def _get_date_key(self) -> str:
        """获取日期键"""
        return datetime.utcnow().strftime("%Y-%m-%d")
    
    def _get_next_reset_time(self) -> float:
        """获取下次重置时间"""
        tomorrow = datetime.utcnow().replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=1)
        return tomorrow.timestamp()
    
    async def get_quota_stats(self, user_id: str = None) -> Dict[str, Any]:
        """获取配额使用统计"""
        if not self.redis:
            return {"error": "Redis not available"}
        
        date_key = self._get_date_key()
        
        # 全局统计
        global_key = f"{self.prefix}:global:{date_key}"
        global_used = await self.redis.get(global_key) or 0
        global_used = int(global_used)
        
        stats = {
            "date": date_key,
            "global": {
                "used": global_used,
                "remaining": max(0, self.daily_quota - global_used),
                "total": self.daily_quota,
                "usage_percent": round(global_used / self.daily_quota * 100, 2)
            },
            "reset_at": self._get_next_reset_time()
        }
        
        # 用户统计
        if user_id:
            user_quota = 1000
            user_key = f"{self.prefix}:user:{user_id}:{date_key}"
            user_used = await self.redis.get(user_key) or 0
            user_used = int(user_used)
            
            stats["user"] = {
                "used": user_used,
                "remaining": max(0, user_quota - user_used),
                "total": user_quota,
                "usage_percent": round(user_used / user_quota * 100, 2)
            }
        
        return stats


# 全局实例
rate_limiter = RateLimiter()
youtube_quota_manager = YouTubeAPIQuotaManager()


def rate_limit(
    max_requests: int = 100,
    window_seconds: int = 60,
    scope: RateLimitScope = RateLimitScope.IP,
    algorithm: RateLimitAlgorithm = RateLimitAlgorithm.FIXED_WINDOW,
    key_func: Callable = None
):
    """
    速率限制装饰器
    
    Args:
        max_requests: 最大请求数
        window_seconds: 时间窗口
        scope: 限流范围
        algorithm: 限流算法
        key_func: 自定义键函数
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # 从参数中提取请求对象
            request = None
            current_user = None
            
            for arg in args:
                if isinstance(arg, Request):
                    request = arg
                elif hasattr(arg, 'id'):  # 假设是用户对象
                    current_user = arg
            
            # 确定标识符
            if key_func:
                identifier = key_func(*args, **kwargs)
            elif scope == RateLimitScope.IP and request:
                identifier = request.client.host
            elif scope == RateLimitScope.USER and current_user:
                identifier = str(current_user.id)
            elif scope == RateLimitScope.ENDPOINT:
                identifier = func.__name__
            else:
                identifier = "global"
            
            # 检查速率限制
            limit_result = await rate_limiter.check_rate_limit(
                scope=scope,
                identifier=identifier,
                max_requests=max_requests,
                window_seconds=window_seconds,
                algorithm=algorithm
            )
            
            if not limit_result["allowed"]:
                logger.warning(f"速率限制触发", 
                             scope=scope.value, 
                             identifier=identifier, 
                             function=func.__name__)
                
                raise HTTPException(
                    status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                    detail="请求过于频繁，请稍后再试",
                    headers={
                        "Retry-After": str(int(limit_result["retry_after"])),
                        "X-RateLimit-Remaining": str(limit_result["remaining"]),
                        "X-RateLimit-Reset": str(int(limit_result["reset_at"])),
                    }
                )
            
            return await func(*args, **kwargs)
        
        return wrapper
    return decorator


def youtube_quota_limit(operation: str):
    """YouTube API配额限制装饰器"""
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # 提取用户ID
            user_id = None
            for arg in args:
                if hasattr(arg, 'id'):  # 假设是用户对象
                    user_id = str(arg.id)
                    break
            
            # 检查配额
            quota_check = await youtube_quota_manager.check_quota(operation, user_id)
            if not quota_check["allowed"]:
                logger.warning(f"YouTube API配额不足", 
                             operation=operation, 
                             user_id=user_id,
                             reason=quota_check.get("reason"))
                
                raise HTTPException(
                    status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                    detail=f"API配额不足: {quota_check.get('reason')}",
                    headers={
                        "X-Quota-Remaining": str(quota_check["remaining"]),
                        "X-Quota-Reset": str(int(quota_check["reset_at"])),
                    }
                )
            
            try:
                # 执行函数
                result = await func(*args, **kwargs)
                
                # 消费配额
                await youtube_quota_manager.consume_quota(operation, user_id)
                
                return result
                
            except Exception as e:
                # 如果函数执行失败，不消费配额
                logger.error(f"函数执行失败，未消费配额: {e}")
                raise
        
        return wrapper
    return decorator


async def get_rate_limit_status(
    scope: RateLimitScope,
    identifier: str,
    max_requests: int,
    window_seconds: int
) -> Dict[str, Any]:
    """获取速率限制状态"""
    return await rate_limiter.check_rate_limit(
        scope=scope,
        identifier=identifier,
        max_requests=max_requests,
        window_seconds=window_seconds
    )